import numpy as np
import matplotlib.pyplot as plt
from functools import reduce
import json


def ivecs_read(fname):
    a = np.fromfile(fname, dtype='int32')
    d = a[0]
    return a.reshape(-1, d + 1)[:, 1:].copy(), d


def get_bias(ip_gnd_idx, l2_gnd_idx, cosine_gnd_idx, norm_gnd_idx):
    n_histogram = 5
    split = 20

    norm_percent_l = np.zeros(shape=n_histogram + 1)
    cosine_percent_l = np.zeros(shape=n_histogram + 1)
    l2_percent_l = np.zeros(shape=n_histogram + 1)

    norm_inter_cosine_percent_l = np.zeros(shape=n_histogram + 1)
    norm_inter_l2_percent_l = np.zeros(shape=n_histogram + 1)
    cosine_inter_l2_percent_l = np.zeros(shape=n_histogram + 1)
    norm_inter_cosine_inter_l2_percent_l = np.zeros(shape=n_histogram + 1)

    n_query = len(ip_gnd_idx)
    ip_topk = len(ip_gnd_idx[0])
    for i in range(n_query):
        if i % 200 == 0:
            print("finish process iter %d " % i)
        l2_split_arr = np.array_split(l2_gnd_idx[i], split)
        cosine_split_arr = np.array_split(cosine_gnd_idx[i], split)
        norm_split_arr = np.array_split(norm_gnd_idx[0], split)

        norm_percent = [len(np.intersect1d(norm_split_arr[j], ip_gnd_idx[i])) / ip_topk
                        for j in range(split)]
        norm_percent[n_histogram] = np.sum(norm_percent[n_histogram:])
        norm_percent = norm_percent[:n_histogram + 1]
        norm_percent_l += norm_percent

        cosine_percent = [len(np.intersect1d(cosine_split_arr[j], ip_gnd_idx[i])) / ip_topk
                          for j in range(split)]
        cosine_percent[n_histogram] = np.sum(cosine_percent[n_histogram:])
        cosine_percent = cosine_percent[:n_histogram + 1]
        cosine_percent_l += cosine_percent

        l2_percent = [len(np.intersect1d(l2_split_arr[j], ip_gnd_idx[i])) / ip_topk
                      for j in range(split)]
        l2_percent[n_histogram] = np.sum(l2_percent[n_histogram:])
        l2_percent = l2_percent[:n_histogram + 1]
        l2_percent_l += l2_percent

        norm_inter_cosine_percent = [
            len(reduce(np.intersect1d, (norm_split_arr[j], cosine_split_arr[j], ip_gnd_idx[i]))) / ip_topk
            for j in range(split)]
        # norm_inter_cosine_percent = [
        #     0 if len(np.intersect1d(np.union1d(norm_split_arr[j], cosine_split_arr[j]), ip_gnd_idx[j])) == 0 else
        #     len(reduce(np.intersect1d, (norm_split_arr[j], cosine_split_arr[j], ip_gnd_idx[j]))) /
        #     len(np.intersect1d(np.union1d(norm_split_arr[j], cosine_split_arr[j]), ip_gnd_idx[j]))
        #     for j in range(split)]
        norm_inter_cosine_percent[n_histogram] = np.sum(norm_inter_cosine_percent[n_histogram:])
        norm_inter_cosine_percent = norm_inter_cosine_percent[:n_histogram + 1]
        norm_inter_cosine_percent_l += norm_inter_cosine_percent

        norm_inter_l2_percent = [
            len(reduce(np.intersect1d, (norm_split_arr[j], l2_split_arr[j], ip_gnd_idx[i]))) / ip_topk
            for j in range(split)]
        # norm_inter_l2_percent = [
        #     0 if len(np.intersect1d(np.union1d(norm_split_arr[j], l2_split_arr[j]), ip_gnd_idx[j])) == 0 else
        #     len(reduce(np.intersect1d, (norm_split_arr[j], l2_split_arr[j], ip_gnd_idx[j]))) /
        #     len(np.intersect1d(np.union1d(norm_split_arr[j], l2_split_arr[j]), ip_gnd_idx[j]))
        #     for j in range(split)]
        norm_inter_l2_percent[n_histogram] = np.sum(norm_inter_l2_percent[n_histogram:])
        norm_inter_l2_percent = norm_inter_l2_percent[:n_histogram + 1]
        norm_inter_l2_percent_l += norm_inter_l2_percent

        cosine_inter_l2_percent = [
            len(reduce(np.intersect1d, (cosine_split_arr[j], l2_split_arr[j], ip_gnd_idx[i]))) / ip_topk
            for j in range(split)]
        # cosine_inter_l2_percent = [
        #     0 if len(np.intersect1d(np.union1d(cosine_split_arr[j], l2_split_arr[j]), ip_gnd_idx[j])) == 0 else
        #     len(reduce(np.intersect1d, (cosine_split_arr[j], l2_split_arr[j], ip_gnd_idx[j]))) /
        #     len(np.intersect1d(np.union1d(cosine_split_arr[j], l2_split_arr[j]), ip_gnd_idx[j]))
        #     for j in range(split)]
        cosine_inter_l2_percent[n_histogram] = np.sum(cosine_inter_l2_percent[n_histogram:])
        cosine_inter_l2_percent = cosine_inter_l2_percent[:n_histogram + 1]
        cosine_inter_l2_percent_l += cosine_inter_l2_percent

        norm_inter_cosine_inter_l2_percent = [
            len(reduce(np.intersect1d,
                       (norm_split_arr[j], cosine_split_arr[j], l2_split_arr[j], ip_gnd_idx[i]))) / ip_topk
            for j in range(split)]
        # norm_inter_cosine_inter_l2_percent = [
        #     0 if len(np.intersect1d(reduce(np.union1d, (norm_split_arr[j], cosine_split_arr[j], l2_split_arr[j])),
        #                             ip_gnd_idx[j])) == 0 else
        #     len(reduce(np.intersect1d, (norm_split_arr[j], cosine_split_arr[j], l2_split_arr[j], ip_gnd_idx[j]))) /
        #     len(np.intersect1d(reduce(np.union1d, (norm_split_arr[j], cosine_split_arr[j], l2_split_arr[j])),
        #                        ip_gnd_idx[j]))
        #     for j in range(split)]
        norm_inter_cosine_inter_l2_percent[n_histogram] = np.sum(norm_inter_cosine_inter_l2_percent[n_histogram:])
        norm_inter_cosine_inter_l2_percent = norm_inter_cosine_inter_l2_percent[:n_histogram + 1]
        norm_inter_cosine_inter_l2_percent_l += norm_inter_cosine_inter_l2_percent

    norm_percent_l /= n_query
    cosine_percent_l /= n_query
    l2_percent_l /= n_query

    norm_inter_cosine_percent_l /= n_query
    norm_inter_l2_percent_l /= n_query
    cosine_inter_l2_percent_l /= n_query
    norm_inter_cosine_inter_l2_percent_l /= n_query
    res_json = {
        'norm': norm_percent_l,
        'cosine': cosine_percent_l,
        'l2': l2_percent_l,
        'norm_inter_cosine': norm_inter_cosine_percent_l,
        'norm_inter_l2': norm_inter_l2_percent_l,
        'cosine_inter_l2': cosine_inter_l2_percent_l,
        'norm_inter_cosine_inter_l2': norm_inter_cosine_inter_l2_percent_l
    }

    return res_json


'''
输入0-5%, 5-10%, 10-15%, 15-20%, 20-25%, 25-100%的百分比, 输出并保存图片
'''


def draw_histogram(percent_value_l, dist_func, ds_name):
    assert len(percent_value_l) == 6
    x = range(len(percent_value_l))
    plt.bar(x, percent_value_l, tick_label=['0-5%', '5-10%', '10-15%', '15-20%', '20-25%', '25-100%'], fc='darkgray')
    plt.xlabel(dist_func)
    plt.ylabel('Percentage')
    # plt.xlabel('Norm Ranking', fontsize=fontsize)
    # plt.ylabel('Percentage', fontsize=fontsize)
    # plt.title('yahoo!music: Position of the ground truth top-10 in norm distribution', fontsize=24)
    plt.title('%s: Position of top-10 in %s distribution' % (ds_name, dist_func))

    plt.xticks()
    plt.yticks()
    # plt.xticks(fontsize=ticksize - 8)
    # plt.yticks(fontsize=ticksize)
    plt.savefig('../topk-norm-distribution/%s_%s.png' % (ds_name, dist_func))
    print("save %s_%s.png" % (ds_name, dist_func))
    plt.close()


if __name__ == '__main__':
    ds_name_l = ['audio', 'imagenet', 'movielens', 'music100', 'netflix', 'normal-64', 'text-to-image', 'tiny5m',
                 'word2vec', 'yahoomusic']
    basic_dir = '/home/zhengbian/mips-graph/data'

    for ds_name in ds_name_l:
        print("start processing %s" % ds_name)
        ip_gnd, _ = ivecs_read('%s/%s_bias/ip_gnd_idx.ivecs' % (basic_dir, ds_name))
        cosine_gnd, _ = ivecs_read('%s/%s_bias/cosine_gnd_idx.ivecs' % (basic_dir, ds_name))
        l2_gnd, _ = ivecs_read('%s/%s_bias/l2_gnd_idx.ivecs' % (basic_dir, ds_name))
        norm_gnd, _ = ivecs_read('%s/%s_bias/norm_gnd_idx.ivecs' % (basic_dir, ds_name))

        percent_l_m = get_bias(ip_gnd, l2_gnd, cosine_gnd, norm_gnd)
        print("norm percent %s" % ds_name, percent_l_m['norm'])
        draw_histogram(percent_l_m['norm'], 'norm', ds_name)
        print("norm cosine %s" % ds_name, percent_l_m['cosine'])
        draw_histogram(percent_l_m['cosine'], 'cosine', ds_name)
        print("norm l2 %s" % ds_name, percent_l_m['l2'])
        draw_histogram(percent_l_m['l2'], 'l2', ds_name)

        print("norm norm inter cosine %s" % ds_name, percent_l_m['norm_inter_cosine'])
        draw_histogram(percent_l_m['norm_inter_cosine'], 'norm_inter_cosine', ds_name)
        print("norm norm inter l2 %s" % ds_name, percent_l_m['norm_inter_l2'])
        draw_histogram(percent_l_m['norm_inter_l2'], 'norm_inter_l2', ds_name)
        print("norm cosine inter l2 %s" % ds_name, percent_l_m['cosine_inter_l2'])
        draw_histogram(percent_l_m['cosine_inter_l2'], 'cosine_inter_l2', ds_name)
        print("norm norm inter cosine inter l2 %s" % ds_name, percent_l_m['norm_inter_cosine_inter_l2'])
        draw_histogram(percent_l_m['norm_inter_cosine_inter_l2'], 'norm_inter_cosine_inter_l2', ds_name)

        with open("../topk-norm-distribution/%s_bias.json" % ds_name, 'w') as f:
            percent_res = {
                'norm': percent_l_m['norm'].tolist(),
                'cosine': percent_l_m['cosine'].tolist(),
                'l2': percent_l_m['l2'].tolist(),
                'norm_inter_cosine': percent_l_m['norm_inter_cosine'].tolist(),
                'norm_inter_l2': percent_l_m['norm_inter_l2'].tolist(),
                'cosine_inter_l2': percent_l_m['cosine_inter_l2'].tolist(),
                'norm_inter_cosine_inter_l2': percent_l_m['norm_inter_cosine_inter_l2'].tolist()
            }
            json.dump(percent_res, f)
